# ==================== 导入需要的工具包 ====================
import torch # 导入PyTorch核心库，提供张量操作和自动微分功能
import torch.nn as nn # 导入神经网络模块，包含Linear、Embedding、LayerNorm等层
import torch.optim as optim # 导入优化器模块，包含Adam、SGD等优化算法
import torch.utils.data as data # 导入数据处理模块，用于DataLoader和Dataset
import math # 导入数学函数库，提供sqrt、log等数学运算
import copy # 导入深度复制模块，用于复制对象而不共享引用


"""
多头注意力机制 - 这是Transformer的核心部分
就像人有多个感官（眼睛、耳朵、鼻子），每个感官关注不同的信息
多头注意力让模型同时关注多个不同的信息
"""
class MultiHeadAttention(nn.Module): # 定义一个类叫多头注意力，继承自nn.Module
    def __init__(self,d_model,num_heads): # 初始化函数，设置参数
        super(MultiHeadAttention,self).__init__() # 调用父类的初始化方法（必须的）
        # assert函数：检查条件，如果为False就抛出异常
        # 检查模型维度能否被头数整除，如果不能就报错
        # 为什么要检查？因为我们要把模型维度平均分配给每个头
        assert d_model % num_heads ==0, "d_model必须能被num_head整除"
        self.d_model=d_model # 保存模型维度（比如512）
        self.num_heads=num_heads # 保存注意力头数（比如8个）
        self.d_k=d_model // num_heads # 计算每个头分到的维度（512÷8=64）

        # nn.Linear函数：创建线性变换层，公式为 y = xW^T + b
        # 创建4个线性变换层，用来把输入变成Q、K、V矩阵
        # 为什么要4个？因为需要分别处理查询(Q)、键(K)、值(V)和输出
        self.W_q=nn.Linear(d_model,d_model) # 查询变换层：把输入变成查询矩阵Q
        self.W_k=nn.Linear(d_model,d_model) # 键变换层：把输入变成键矩阵K
        self.W_v=nn.Linear(d_model,d_model) # 值变换层：把输入变成值矩阵V
        self.W_o=nn.Linear(d_model,d_model) # 输出变换层：把多头结果合并成最终输出

    def scaled_dot_product_attention(self,Q,K,V,mask=None): # 计算注意力的核心函数
        """
        这个函数计算注意力权重
        就像人看东西时，眼睛会重点关注某些部分，忽略其他部分
        """
        # torch.matmul函数：矩阵乘法运算
        # K.transpose(-2,-1)：转置K矩阵的最后两个维度
        # math.sqrt函数：计算平方根
        # 计算Q和K的相似度分数，然后除以sqrt(d_k)进行缩放
        # 为什么要缩放？防止分数太大导致梯度消失
        attn_scores=torch.matmul(Q,K.transpose(-2,-1)) / math.sqrt(self.d_k)

        # if语句：条件判断
        # masked_fill函数：用指定值填充掩码为True的位置
        # 如果有掩码，就把不需要的位置设为很小的数
        # 为什么要掩码？防止模型看到不应该看到的信息（比如未来的词）
        if mask is not None:
            attn_scores=attn_scores.masked_fill(mask == 0,-1e9)
        
        # torch.softmax函数：softmax激活函数，将输入转换为概率分布
        # dim=-1：在最后一个维度上应用softmax
        # 用softmax把分数变成概率（所有概率加起来等于1）
        # 为什么要用softmax？让模型知道哪些位置最重要
        attn_probs =torch.softmax(attn_scores,dim=-1)

        # 用注意力权重对V进行加权求和，得到最终输出
        # 为什么要加权求和？让重要的信息影响更大
        output=torch.matmul(attn_probs,V)
        return output # 返回注意力机制的输出结果
    
    def split_heads(self,x): # 把输入分成多个头
        """
        就像把一个大蛋糕切成8块，每块大小一样
        为什么要分头？让每个头关注不同类型的信息
        """
        # x.size()函数：获取张量的形状
        batch_size,seq_length,d_model = x.size() # 获取输入的大小
        # x.view()函数：改变张量的形状，不改变数据
        # transpose(1,2)函数：交换第1和第2个维度
        # 把输入重新排列，分成多个头
        return x.view(batch_size,seq_length,self.num_heads,self.d_k).transpose(1,2)
    
    def combine_head(self,x): # 把多个头的输出合并
        """
        就像把8块蛋糕重新拼成一个大蛋糕
        为什么要合并？把所有头学到的信息综合起来
        """
        # 获取输入张量的形状，_表示忽略某个维度
        batch_size,_,seq_length,d_k=x.size() # 获取输入的大小
        # contiguous()函数：确保张量在内存中是连续的
        # 把多个头的输出合并成一个输出
        return x.transpose(1,2).contiguous().view(batch_size,seq_length,self.d_model)
    
    def forward(self,Q,K,V,mask=None): # 前向传播函数
        """
        这是整个多头注意力的主要流程
        就像工厂的流水线，一步一步处理输入
        """
        # 第一步：把输入变成Q、K、V并分成多个头
        Q=self.split_heads(self.W_q(Q)) # 处理查询：线性变换+分头
        K=self.split_heads(self.W_k(K)) # 处理键：线性变换+分头
        V=self.split_heads(self.W_v(V)) # 处理值：线性变换+分头

        # 第二步：计算注意力
        attn_output=self.scaled_dot_product_attention(Q,K,V,mask)

        # 第三步：合并多头并输出
        output=self.W_o(self.combine_head(attn_output))
        return output

"""
前馈神经网络 - 对每个位置独立处理
就像对每个词都做一次"思考"，不依赖其他词
"""
class PositionWiseFeedForward(nn.Module): # 定义前馈网络类
    def __init__(self,d_model,d_ff): # 初始化函数
        super(PositionWiseFeedForward,self).__init__() # 调用父类初始化
        # 创建两个全连接层，先扩大再缩小
        # 为什么要先扩大再缩小？增加模型的表达能力
        self.fc1=nn.Linear(d_model,d_ff) # 第一层：把维度扩大（比如512→2048）
        self.fc2=nn.Linear(d_ff,d_model) # 第二层：把维度缩小回原来的大小（2048→512）
        self.relu=nn.ReLU() # 创建ReLU激活函数，公式为 max(0, x)
    def forward(self,x): # 前向传播
        # 依次通过两个全连接层和激活函数
        return self.fc2(self.relu(self.fc1(x)))

"""
位置编码 - 告诉模型每个词在句子中的位置
因为Transformer没有循环结构，所以需要额外告诉它位置信息
"""
class PositionalEncoding(nn.Module): # 定义位置编码类
    def __init__(self,d_model,max_seq_length): # 初始化函数
        super(PositionalEncoding,self).__init__() # 调用父类初始化
        # torch.zeros函数：创建全零张量
        pe =torch.zeros(max_seq_length,d_model) # 创建位置编码矩阵，先全部填0
        # torch.arange函数：创建等差数列
        # unsqueeze(1)函数：在第1个位置插入维度1
        # 创建位置索引，从0到最大序列长度
        position = torch.arange(0,max_seq_length,dtype=torch.float).unsqueeze(1)
        # torch.exp函数：计算指数
        # torch.log函数：计算对数
        # 计算用于不同频率的除数
        div_term = torch.exp(torch.arange(0,d_model,2).float()*-(math.log(10000.0) /d_model))
        # 切片操作[:,0::2]：取所有行，从第0列开始每隔2列取一列
        # torch.sin函数：计算正弦值
        # 偶数位置用正弦函数
        pe[:,0::2] =torch.sin(position*div_term)
        # 切片操作[:,1::2]：取所有行，从第1列开始每隔2列取一列
        # torch.cos函数：计算余弦值
        # 奇数位置用余弦函数
        pe[:,1::2] =torch.cos(position*div_term)
        # register_buffer函数：将张量注册为缓冲区，不参与梯度计算
        # unsqueeze(0)函数：在第0个位置插入维度1
        # 把位置编码保存起来，不参与训练
        self.register_buffer('pe',pe.unsqueeze(0))
    def forward(self,x): # 前向传播
        # x.size(1)函数：获取第1个维度的大小（序列长度）
        # 把位置编码加到输入上
        return x+self.pe[:,:x.size(1)]
    
"""
编码器层 - 理解输入序列
就像阅读理解，把输入的文章理解一遍
"""
class EncoderLayer(nn.Module): # 定义编码器层类
    def __init__(self,d_model,num_heads,d_ff,dropout): # 初始化函数
        super(EncoderLayer,self).__init__() # 调用父类初始化
        # 创建自注意力机制，让模型关注输入序列内部的关系
        self.self_attn=MultiHeadAttention(d_model,num_heads)
        # 创建前馈网络，对每个位置进行非线性变换
        self.feed_forward=PositionWiseFeedForward(d_model,d_ff)
        # nn.LayerNorm函数：创建层归一化，对每个样本的每个特征进行归一化
        # 创建层归一化，让训练更稳定
        self.norm1=nn.LayerNorm(d_model) # 第一个归一化层
        self.norm2=nn.LayerNorm(d_model) # 第二个归一化层
        # nn.Dropout函数：创建dropout层，训练时随机将输入设为0
        # 创建dropout，防止过拟合（随机关闭一些神经元）
        self.dropout=nn.Dropout(dropout)
    
    def forward(self,x,mask): # 前向传播函数
        # 第一步：自注意力机制
        # 让模型关注输入序列中每个词与其他词的关系
        attn_output=self.self_attn(x,x,x,mask) # Q、K、V都来自同一个输入
        # 残差连接+层归一化：原始输入+注意力输出
        # 为什么要残差连接？让梯度更容易传播，训练更稳定
        x=self.norm1(x+self.dropout(attn_output))

        # 第二步：前馈网络
        # 对每个位置独立进行非线性变换
        ff_output=self.feed_forward(x)
        # 残差连接+层归一化：原始输入+前馈输出
        x =self.norm2(x+self.dropout(ff_output))
        return x

"""
解码器层 - 生成输出序列
就像写作文，根据理解的内容生成新的句子
"""
class DecoderLayer(nn.Module): # 定义解码器层类
    def __init__(self,d_model,num_heads,d_ff,dropout): # 初始化函数
        super(DecoderLayer,self).__init__() # 调用父类初始化
        # 创建自注意力机制，让模型关注目标序列内部的关系
        self.self_attn=MultiHeadAttention(d_model,num_heads)
        # 创建交叉注意力机制，让模型关注编码器的输出
        self.cross_attn=MultiHeadAttention(d_model,num_heads)
        # 创建前馈网络
        self.feed_forward=PositionWiseFeedForward(d_model,d_ff)
        # 创建三个层归一化
        self.norm1=nn.LayerNorm(d_model) # 自注意力后的归一化
        self.norm2=nn.LayerNorm(d_model) # 交叉注意力后的归一化
        self.norm3=nn.LayerNorm(d_model) # 前馈网络后的归一化
        # 创建dropout
        self.dropout=nn.Dropout(dropout)

    def forward(self,x,enc_output,src_mask,tgt_mask): # 前向传播函数
        # 第一步：自注意力机制
        # 让模型关注目标序列中每个词与其他词的关系
        attn_output=self.self_attn(x,x,x,tgt_mask)
        x=self.norm1(x+self.dropout(attn_output)) # 残差连接+归一化
        
        # 第二步：交叉注意力机制
        # 让模型关注编码器的输出，理解源序列的信息
        cross_attn_output=self.cross_attn(x,enc_output,enc_output,src_mask)
        x=self.norm2(x+self.dropout(cross_attn_output)) # 残差连接+归一化
        
        # 第三步：前馈网络
        # 对每个位置进行非线性变换
        ff_output=self.feed_forward(x)
        x=self.norm3(x+self.dropout(ff_output)) # 残差连接+归一化
        return x

"""
完整的Transformer模型
就像一台翻译机器，输入一种语言，输出另一种语言
"""
class Transformer(nn.Module): # 定义Transformer类
    def __init__(self,src_vocab_size,tgt_vocab_size,d_model,num_heads,num_layers,d_ff,max_seq_length,dropout): # 初始化函数
        super(Transformer,self).__init__() # 调用父类初始化
        # nn.Embedding函数：创建词嵌入层，将词汇索引转换为密集向量
        # 创建词嵌入层，把词变成数字向量
        self.encoder_embedding=nn.Embedding(src_vocab_size,d_model) # 源语言词嵌入
        self.decoder_embedding=nn.Embedding(tgt_vocab_size,d_model) # 目标语言词嵌入
        # 创建位置编码层
        self.positional_encoding=PositionalEncoding(d_model,max_seq_length)
        # nn.ModuleList函数：创建模块列表，用于存储多个子模块
        # 创建多个编码器层，一层一层地理解输入
        self.encoder_layers=nn.ModuleList([EncoderLayer(d_model,num_heads,d_ff,dropout) for _ in range(num_layers)])
        # 创建多个解码器层，一层一层地生成输出
        self.decoder_layers=nn.ModuleList([DecoderLayer(d_model,num_heads,d_ff,dropout) for _ in range(num_layers)])
        # 创建最终输出层，把模型输出变成词汇表上的概率
        self.fc= nn.Linear(d_model,tgt_vocab_size)
        # 创建dropout层
        self.dropout=nn.Dropout(dropout)
    
    def generate_mask(self,src,tgt): # 生成掩码函数
        # != 操作符：不等于比较
        # unsqueeze函数：在指定位置插入维度1
        # 源序列掩码：标记哪些位置是真实的词，哪些是填充
        # 为什么要掩码？防止模型关注填充位置
        src_mask=(src!=0).unsqueeze(1).unsqueeze(2)
        
        # 目标序列掩码：标记哪些位置是真实的词
        tgt_mask=(tgt!=0).unsqueeze(1).unsqueeze(3)
        # size(1)函数：获取第1个维度的大小
        seq_length=tgt.size(1) # 获取序列长度
        # torch.ones函数：创建全1张量
        # torch.triu函数：创建上三角矩阵
        # diagonal=1：从主对角线+1开始
        # 创建上三角掩码，防止模型看到未来的词
        # 为什么要这个掩码？防止作弊（在翻译时不能看到后面的词）
        nopeak_mask=(1-torch.triu(torch.ones(1,seq_length,seq_length),diagonal=1)).bool()
        # & 操作符：逻辑与运算
        tgt_mask=tgt_mask & nopeak_mask # 把两个掩码合并
        return src_mask,tgt_mask

    def forward(self,src,tgt): # 前向传播函数
        # 生成掩码
        src_mask,tgt_mask=self.generate_mask(src,tgt)

        # 编码器部分：理解源序列
        # 把词变成向量，加上位置信息，然后通过多个编码器层
        src_embedded=self.dropout(self.positional_encoding(self.encoder_embedding(src)))
        enc_output=src_embedded # 初始化编码器输出
        # for循环：遍历编码器层列表
        for enc_layer in self.encoder_layers: # 通过每个编码器层
            enc_output =enc_layer(enc_output,src_mask)

        # 解码器部分：生成目标序列
        # 把目标词变成向量，加上位置信息，然后通过多个解码器层
        tgt_embedded=self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
        dec_output =tgt_embedded # 初始化解码器输出
        # 遍历解码器层列表
        for dec_layer in self.decoder_layers: # 通过每个解码器层
            dec_output=dec_layer(dec_output,enc_output,src_mask,tgt_mask)
        
        # 最终输出：把解码器输出变成词汇表上的概率
        output=self.fc(dec_output)
        return output

# ==================== 设置模型参数 ====================
src_vocab_size=5000 # 源语言有多少个不同的词
tgt_vocab_size=5000 # 目标语言有多少个不同的词
d_model=512 # 词向量的维度（每个词用512个数字表示）
num_heads=8 # 注意力头数（8个并行的注意力）
num_layers=6 # 编码器和解码器各有6层
d_ff=2048 # 前馈网络的隐藏层维度
max_seq_length=100 # 最大句子长度
dropout=0.1 # 训练时随机关闭10%的神经元

# ==================== 创建模型 ====================
transformer=Transformer(src_vocab_size,tgt_vocab_size,d_model,num_heads,num_layers,d_ff,max_seq_length,dropout)

# ==================== 准备训练数据 ====================
# torch.randint函数：生成随机整数张量
# 生成一些随机数据来训练模型
src_data=torch.randint(1,src_vocab_size,(64,max_seq_length)) # 源序列数据
tgt_data=torch.randint(1,tgt_vocab_size,(64,max_seq_length)) # 目标序列数据

# ==================== 设置训练工具 ====================
# nn.CrossEntropyLoss函数：创建交叉熵损失函数
# ignore_index=0：忽略索引为0的标签（填充位置）
# 定义损失函数：计算预测和真实答案的差距
criterion=nn.CrossEntropyLoss(ignore_index=0) # 忽略填充位置的损失
# optim.Adam函数：创建Adam优化器
# lr：学习率，控制参数更新的步长
# betas：动量参数，控制梯度的指数移动平均
# eps：数值稳定性参数
# 定义优化器：用来更新模型参数，让模型学习
optimizer=optim.Adam(transformer.parameters(),lr=0.0001,betas=(0.9,0.98),eps=1e-9)

# ==================== 开始训练 ====================
# train()函数：设置模型为训练模式，启用dropout和batch normalization
transformer.train() # 设置模型为训练模式
# range函数：生成数字序列
for epoch in range(100): # 训练100轮
    # zero_grad()函数：将梯度清零，防止梯度累积
    optimizer.zero_grad() # 清空梯度（每次训练前都要清空）

    # 前向传播：让模型预测
    # [:,:-1]切片：取所有行，去掉最后一列
    output=transformer(src_data,tgt_data[:,:-1]) # 输入源序列和目标序列（去掉最后一个词）
    
    # 计算损失：比较预测结果和真实答案
    # contiguous()函数：确保张量在内存中连续
    # view()函数：改变张量形状
    # [:,1:]切片：取所有行，从第1列开始（去掉第0列）
    loss=criterion(
        output.contiguous().view(-1,tgt_vocab_size), # 把输出变成二维
        tgt_data[:,1:].contiguous().view(-1) # 把目标变成一维
    )
    # backward()函数：反向传播，计算梯度
    loss.backward() # 反向传播：计算梯度
    # step()函数：更新模型参数
    optimizer.step() # 更新模型参数
    # f-string：格式化字符串，{变量}会被替换为实际值
    # item()函数：将标量张量转换为Python数值
    # :.4f：保留4位小数的浮点数格式
    print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}') # 打印训练进度